package com.rmfyzxfw.caseprocess.util;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.pool.DruidPooledConnection;
import com.alibaba.druid.spring.boot.autoconfigure.DruidDataSourceBuilder;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.rmfyzxfw.caseprocess.config.datasource.DataSourceContextHolder;
import com.rmfyzxfw.caseprocess.config.datasource.DataSourceRouting;
import com.rmfyzxfw.caseprocess.enums.DynamicDatabaseTypeEnum;
import com.rmfyzxfw.caseprocess.mapper.DynamicDataSourceMapper;
import com.rmfyzxfw.caseprocess.model.po.DynamicDataSource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.Serializable;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;

@Component
public class DynamicDataSourceUtil {

    private static DataSourceRouting dataSourceRouting;
    private static DynamicDataSourceMapper dataSourceMapper;
    private final static Map<Serializable, DynamicDataSource> dataSourceMap = new ConcurrentHashMap<>();

    @Autowired
    public void setDynamicDataSourceRepository(DynamicDataSourceMapper dataSourceMapper) {
        DynamicDataSourceUtil.dataSourceMapper = dataSourceMapper;
    }

    @Autowired(required = false)
    public void setDataSourceRouting(DataSourceRouting dataSourceRouting) {
        DynamicDataSourceUtil.dataSourceRouting = dataSourceRouting;
    }

    public static <T> T route(Serializable dataSourceId, Supplier<T> supplier) {
        if (dataSourceId == null) {
            return supplier.get();
        }

        try {
            String dataSourceKey = "outerDataSource-" + dataSourceId;
            DynamicDataSource dynamicDataSource = dataSourceMap.get(dataSourceId);
            if(dynamicDataSource == null){
                dynamicDataSource = dataSourceMapper.selectOne(Wrappers.<DynamicDataSource>lambdaQuery().eq(DynamicDataSource::getDynamicDataSourceId, dataSourceId));
                dataSourceMap.put(dataSourceId, dynamicDataSource);
            }
            addDataSourceIfNotExist(dataSourceKey, dynamicDataSource);
            DataSourceContextHolder.setDataSourceKey(dataSourceKey);
            DataSourceContextHolder.setDataSource(dynamicDataSource);
            T t = supplier.get();
            return t;
        } finally {
            DataSourceContextHolder.clearDataSourceKey();
            DataSourceContextHolder.clearDataSource();
        }

    }


    private static void addDataSourceIfNotExist(String dataSourceType, DynamicDataSource dynamicDataSource) {
        if (dynamicDataSource != null && !dataSourceRouting.outerDataSourceContainsKey(dataSourceType)) {
            //创建数据源
            DruidDataSource dataSource = DruidDataSourceBuilder.create().build();
            dataSource.setName(dataSourceType);
            dataSource.setUrl(dynamicDataSource.getUrl());
            dataSource.setUsername(dynamicDataSource.getUsername());
            dataSource.setPassword(dynamicDataSource.getPassword());
            dataSource.setConnectionErrorRetryAttempts(3);
            dataSource.setBreakAfterAcquireFailure(true);
            dataSource.setMaxWait(5000L);
            dataSource.setTestOnBorrow(true);
            dataSource.setDriverClassName(DynamicDatabaseTypeEnum.ofCode(dynamicDataSource.getDatabaseType()).getDriver());
            testConnection(dataSource);
            dataSourceRouting.addOuterDataSource(dataSourceType, dataSource);
        }
    }

    public static void removeDataSourceIfExist(Long dataSourceId) {
        if (dataSourceId == null) {
            return;
        }
        String dataSourceKey = "outerDataSource-" + dataSourceId;
        dataSourceRouting.removeOuterDataSource(dataSourceKey);
    }

    public static void testConnection(DruidDataSource dataSource) {
        if (dataSource == null) {
            throw new RuntimeException("数据源不存在");
        }
        try (DruidPooledConnection connection = dataSource.getConnection(3000L)) {
            if (connection == null) {
                throw new RuntimeException("连接不存在");
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
